import os
import uuid
from typing import Union
import pandas as pd
from datasets import load_dataset, Dataset, load_from_disk


def binary_result_parse(dataset: Dataset, output_filename, only_true):
    """
    def get_last_folder(path):
    return os.path.basename(path.rstrip('/'))

    kws = ["半导体", "微电子", "电子元器件", "芯片", "设备和材料", "集成电路"]

    for kw in kws:
        folder = f"/home/jie/github/LLM/API/openai_examples/industry_dataset/data/csvs/宁波_{kw}"

        result_parse(
            folder,
            os.path.join(
                "/home/jie/github/LLM/API/openai_examples/industry_dataset/output",
                f"{get_last_folder(folder)}.csv"
            ),
        )
    """
    if isinstance(dataset, str):
        assert os.path.exists(dataset)
        dataset = load_from_disk(dataset)

    results = []

    for item in dataset:
        infer = item["infer"]
        rear = infer[-3:]

        if "是" in rear and "否" not in rear:
            cls_target = "是"
        elif "否" in rear and "是" not in rear:
            cls_target = "否"
        else:
            cls_target = rear
        results.append(cls_target)

    new_dataset = dataset.add_column("result", results)

    if only_true:
        new_dataset = new_dataset.filter(lambda x: x["result"] == "是")

    new_dataset.to_csv(output_filename)


def binary_parse_list(data: list, rear_len=-4) -> list:
    results = []
    for infer in data:
        rear = infer[rear_len:]

        if "是" in rear and "否" not in rear:
            cls_target = "是"
        elif "否" in rear and "是" not in rear:
            cls_target = "否"
        else:
            cls_target = infer
        results.append(cls_target)
    return results
